#include "Compiler.h"

namespace FlooP
{
	void Compiler::Compile (Environment& env, std::shared_ptr<AST> ast)
	{
		Compiler compiler = Compiler(env);

		// Compiler::visitor pattern
		compiler.visit (*ast);
	}

	Compiler::Compiler (Environment& env)
	:
		m_env (env)
	{}

	void Compiler::visit (const AST& ast)
	{
		const std::vector<std::shared_ptr<ASTProcedure>>& procedure = ast.procedure ();

		for (uint32_t i = 0; i < procedure.size (); i++) {
			procedure[i]->accept (*this);
		}
	}

	void Compiler::visit (const ASTProcedure& procedure)
	{
		std::shared_ptr<Procedure> new_procedure (new Procedure (procedure.name (), procedure.args ()));
		m_env.add_procedure (new_procedure);

		for (int32_t i = procedure.args ().size () - 1; i >= 0; i--) {
			new_procedure->code ()->write (OP_SET, std::shared_ptr<Object> (new String (procedure.args ()[i])));
			new_procedure->code ()->write (OP_POP);
		}

		if (procedure.body ().get () != nullptr) {
			procedure.body ()->accept (*this, new_procedure);

			new_procedure->code ()->write (OP_RETURN);
		}

		new_procedure->fix_abort_quit ();
	}

	void Compiler::visit (const ASTBlock& block, std::shared_ptr<Procedure> func)
	{
		func->add_block (block.block_num ());

		const std::vector<std::shared_ptr<ASTStatement>>& stmts = block.stmts ();

		for (uint32_t i = 0; i < stmts.size (); i++) {
			stmts[i]->accept (*this, func);
		}

		func->set_block_end (block.block_num (), func->code ()->count ());
	}

	void Compiler::visit (const ASTIf& if_stmt, std::shared_ptr<Procedure> func)
	{
		std::shared_ptr<CodeBlock> code = func->code ();

		// Compile the condition
		if_stmt.cond ()->accept (*this, func);

		uint32_t index = code->count ();

		// If its false branch past the if part
		code->write (OP_BRF, 0);

		// Compile the if part
		if_stmt.if_part ()->accept (*this, func);

		// We set the offset after we compile the if part
		Instruction& brf = (*code)[index];
		brf.m_arg = code->count ();
	}

	void Compiler::visit (const ASTLoop& loop, std::shared_ptr<Procedure> func)
	{
		std::shared_ptr<CodeBlock> code = func->code ();

		// Push our counter onto the stack
		code->write (OP_PUSH, std::shared_ptr<Number> (new Number (0)));

		uint32_t index_start = code->count ();

		// Copy the counter and test against the upper bound
		code->write (OP_COPY);
		loop.times ()->accept (*this, func);
		code->write (OP_EQ);

		// If we have reached our looping limit so branch past the loop
		code->write (OP_BRT, 0);
		uint32_t index_branch = code->count () - 1;

		// If not we execute the loop
		loop.body ()->accept (*this, func);

		// Increment counter by 1 and branch back
		code->write (OP_PUSH, std::shared_ptr<Number> (new Number (1)));
		code->write (OP_ADD);
		code->write (OP_BR, index_start);

		// We pop the counter from the stack
		code->write (OP_POP);

		// Write the branch after the loop is compiled
		Instruction& brf = (*code)[index_branch];
		brf.m_arg = code->count ();
	}

	void Compiler::visit (const ASTMuLoop& loop, std::shared_ptr<Procedure> func)
	{
		std::shared_ptr<CodeBlock> code = func->code ();

		uint32_t index_start = code->count ();

		// We execute the loop
		loop.body ()->accept (*this, func);

		// We write NOP's here
		code->write (OP_NOP);
		code->write (OP_BR, index_start);

		// We pop the counter from the stack
		code->write (OP_POP);
	}

	void Compiler::visit (const ASTAbort& abort, std::shared_ptr<Procedure> func)
	{
		// We save the block number in the branch argument to be fixed later
		func->code ()->write (OP_ABORT, abort.block ());
	}

	void Compiler::visit (const ASTQuit& quit, std::shared_ptr<Procedure> func)
	{
		// We save the block number in the branch argument to be fixed later
		func->code ()->write (OP_QUIT, quit.block ());
	}

	void Compiler::visit (const ASTAssignment& assign, std::shared_ptr<Procedure> func)
	{
		assign.expr ()->accept (*this, func);

		uint32_t cell_num = assign.cell ();

		// Remember we reserved cell 0 for OUTPUT
		if (cell_num == 0) {
			func->code ()->write (OP_SOUTPUT);
		} else {
			func->code ()->write (OP_SCELL, cell_num--);
		}

		func->code ()->write (OP_POP);
	}

	void Compiler::visit (const ASTCompare& expr, std::shared_ptr<Procedure> func)
	{
		expr.left ()->accept (*this, func);
		expr.right ()->accept (*this, func);

		std::shared_ptr<CodeBlock> code = func->code ();
		
		switch (expr.op_type ())
		{
			case ASTCompare::OR:		   code->write (OP_OR);  break;
			case ASTCompare::AND:		   code->write (OP_AND); break;
			case ASTCompare::EQUALS:	   code->write (OP_EQ);  break;
			case ASTCompare::LESS_THAN:	   code->write (OP_LT);  break;
			case ASTCompare::GREATER_THAN: code->write (OP_GT);  break;
		}
	}

	void Compiler::visit (const ASTBinaryMathOp& expr, std::shared_ptr<Procedure> func)
	{
		expr.left ()->accept (*this, func);
		expr.right ()->accept (*this, func);

		std::shared_ptr<CodeBlock> code = func->code ();
		
		switch (expr.op_type ())
		{
			case ASTBinaryMathOp::PLUS: code->write (OP_ADD); break;
			case ASTBinaryMathOp::STAR: code->write (OP_MUL); break;
		}
	}

	void Compiler::visit (const ASTCell& expr, std::shared_ptr<Procedure> func)
	{
		uint32_t cell_num = expr.value ();

		// Remember we reserved cell 0 for OUTPUT
		if (cell_num == 0) {
			func->code ()->write (OP_ROUTPUT);
		} else {
			func->code ()->write (OP_RCELL, cell_num--);
		}
	}

	void Compiler::visit (const ASTProcedureCall& expr, std::shared_ptr<Procedure> func)
	{
		std::shared_ptr<CodeBlock> code = func->code ();
		const std::vector<std::shared_ptr<ASTExpression>>& args = expr.args ();

		std::shared_ptr<Procedure> call_procedure = m_env.find_procedure (expr.name ());

		if (call_procedure == nullptr) {
			throw Error ("Compiler : Procedure '%s' has not been defined.", expr.name ());
		}

		if (args.size () != call_procedure->args ().size ()) {
			throw Error ("Compiler : Procedure '%s' expects %i arguments.", expr.name ().c_str (), call_procedure->args ().size ());
		}

		for (uint32_t i = 0; i < args.size (); i++) {
			args[i]->accept (*this, func);
		}

		code->write (OP_CALL, std::shared_ptr<Object> (new String (expr.name ())));
	}

	void Compiler::visit (const ASTIdentifier& expr, std::shared_ptr<Procedure> func)
	{
		std::shared_ptr<Object> var = func->scope ()->find (expr.name ());

		if (var.get () == nullptr) {
			throw Error ("Compiler : Undefined procedure variable '%s'.", expr.name ().c_str ());
		}

		func->code ()->write (OP_REF, std::shared_ptr<Object> (new String (expr.name ())));
	}

	void Compiler::visit (const ASTNumber& num, std::shared_ptr<Procedure> func)
	{
		func->code ()->write (OP_PUSH, std::shared_ptr<Object> (new Number (num.value ())));
	}

	void Compiler::visit (const ASTString& str, std::shared_ptr<Procedure> func)
	{
		func->code ()->write (OP_PUSH, std::shared_ptr<Object> (new String (str.text ())));
	}
}